from . import custommodule
from datetime import datetime, timedelta
import sqlite3
import pandas as pd
import logging

# 配置日志记录器，记录回测过程中买卖操作信息
logging.basicConfig(filename='operate.log', filemode='a', level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S', encoding='utf-8')


# 交易策略回测（默认策略）
def backtest(strategy=None, stock_codes=None, start_date=None, end_date=None):
    """
    参数：
    choice：选择的操作策略
    stock_codes：股票池，默认为None
    start_date：回测起始日期，默认为None
    end_date：回测结束日期，默认为None
    """
    strategys = custommodule.Strategy.strategys()
    # 检查输入的交易策略参数
    while True:
        if strategy is None:
            for x, y in strategys.items():
                print(f"{x}. {y}")
            strategy = input("请输入回测策略（默认策略）：")
        if strategy and int(strategy) in strategys.keys():
            strategy = int(strategy)
            break
        else:
            if strategy == "":
                strategy = 1
            else:
                strategy = None
                print('您的交易策略输入有误，请重新输入')
        
    # 先清空该策略原有的回测结果
    with sqlite3.connect('db_file.db') as conn:
        cursor = conn.cursor()
        # 检查 backtest 表是否存在
        cursor.execute("PRAGMA table_info(backtest)")
        table_info = cursor.fetchall()
        if table_info:
            # 表存在，执行更新操作
            cursor.execute(f"UPDATE backtest SET {strategys[strategy]} = NULL")
        conn.commit()

    # 检查输入的股票起始日期参数
    if start_date is None:start_date = input("请输入回测起始日期（默认为2020-01-01）：")
    if not start_date: start_date = "2020-01-01"

    # 检查输入的股票结束日期参数
    if end_date is None:end_date = input("请输入回测结束日期（默认为当前日期）：")
    if not end_date:end_date = datetime.now().strftime("%Y-%m-%d")

    # 检查输入的股票池参数
    if stock_codes is None:stock_codes = [code[0] for code in custommodule.get_watchlist()]
    # 去除股票代码后缀
    stock_codes = custommodule.StockCodeConversion.remove(stock_codes)
    
    # 开始逐股回测
    for stock_code in stock_codes:
        capital = 1000000 # 原始资金
        stock_num = 0 # 每次交易数量
        stock_nums = 0 # 持仓数量
        cost_price = 0 # 持仓成本价格
        stop_price = 0 # 止损价
        position = 0  # 当前仓位（0：空仓，1：40%仓位，2：70%仓位，3：100%仓位）
        buy_fee = 0.003 # 购买手续费
        sell_fee = 0.0035 # 卖出手续费

        # 按策略计算该股每天的操作建议
        operate = custommodule.Strategy(strategy=strategy, stock_codes=stock_code, start_date=start_date, end_date=end_date).result
        # 如果该股没有操作建议，则跳过
        if operate is None or operate.empty:continue
        
        # 获取该股的交易日期列表
        stock_dates = operate['trade_date'].tolist()
        operate = operate.set_index(['stock_code', 'trade_date'])

        # 开始按日回测
        print(stock_code, "回测中")
        for stock_date in stock_dates:
            # 获取该日期的操作建议
            operating = operate.loc[(stock_code, stock_date), 'operation']
            # 获取该日期的股票价格
            stock_price = custommodule.get_stock_price(stock_code, stock_date)
            # 获取该日期上个交易日的股票价格（用于计算止损价）
            yesterday_price = custommodule.get_previous_price(stock_code, stock_date)
            if yesterday_price is None:yesterday_price = stock_price # 如果上个交易日没有股票价格，则用当前价格代替
            # 获取该日期下个交易日的开盘价（用于交易）
            next_price = custommodule.get_next_price(stock_code, stock_date)
            if next_price is None:next_price = stock_price # 如果下个交易日没有股票价格，则用当前价格代替

            # 是否进行止盈操作(股价上涨10%且仓位大于1则进行卖出操作，仓位为1不进行止盈卖出)
            if (stock_price > cost_price * 1.1) and (position > 1):
                operating = -1 # 止盈卖出
                cost_price = stock_price #用当前价格作为新的成本价

            # 是否进行止损操作
            if stock_price < stop_price:operating = -1 # 止损卖出
            # 计算最新止损价
            stop_price = stock_price * 0.9 if stock_price > yesterday_price else yesterday_price * 0.9

            if operating == 1:
                # 有效性检查
                if stock_price <= 0:
                    continue  # 跳过无效价格的买入操作
                # 使用检查后的有效价格计算
                match position:
                    case 0:
                        stock_num = 400000 // (next_price * (1 + buy_fee))
                        capital = round(capital - stock_num * next_price * (1 + buy_fee), 2)
                        stock_nums = stock_num
                        cost_price = next_price # 记录持仓成本价格
                        position = 1  # 40%仓位
                    case 1:
                        stock_num = (capital / 2) // (next_price * (1 + buy_fee))
                        capital = round(capital - stock_num * next_price * (1 + buy_fee), 2)
                        stock_nums = stock_nums + stock_num
                        cost_price = round((cost_price * (stock_nums - stock_num) + next_price * stock_num) / stock_nums, 2) # 记录持仓成本价格
                        position = 2  # 70%仓位
                    case 2:
                        stock_num = capital // (next_price * (1 + buy_fee))
                        capital = round(capital - stock_num * next_price * (1 + buy_fee), 2)
                        stock_nums = stock_nums + stock_num
                        cost_price = round((cost_price * (stock_nums - stock_num) + next_price * stock_num) / stock_nums, 2) # 记录持仓成本价格
                        position = 3  # 100%仓位
                    case 3:
                        continue  # 100%仓位时，不再买入

                logging.info(
                    f"{stock_code} {stock_date}  买入价格：{float(next_price):>7.2f}  买入数量：{stock_num:>12.2f}  剩余资金：{capital:>12.2f}  剩余股票：{stock_nums}")

            elif operating == -1:
                # 有效性检查
                if stock_price <= 0:
                    continue  # 跳过无效价格的卖出操作
                match position:
                    case 0:
                        continue  # 空仓时，不再卖出
                    case 1:
                        stock_num = stock_nums
                        capital = round(capital + stock_num * next_price * (1 - sell_fee), 2)
                        stock_nums = 0
                        position = 0  # 仓位变为空仓
                    case 2:
                        stock_num = stock_nums * 3 // 7
                        capital = round(capital + stock_num * next_price * (1 - sell_fee), 2)
                        stock_nums = stock_nums - stock_num
                        position = 1  # 仓位变为40%仓位
                    case 3:
                        stock_num = stock_nums * 3 // 10
                        capital = round(capital + stock_num * next_price * (1 - sell_fee), 2)
                        stock_nums = stock_nums - stock_num
                        position = 2  # 仓位变为70%仓位

                logging.info(
                    f"{stock_code} {stock_date}  卖出价格：{float(next_price):>7.2f}  卖出数量：{stock_num:>12.2f}  剩余资金：{capital:>12.2f}  剩余股票：{stock_nums}")

        RoR = round((capital + stock_nums * custommodule.get_stock_price(stock_code, stock_dates[-1]) - 1000000) / 1000000 * 100, 2)
        print(stock_code, "回测结束：剩余资金：", capital, "剩余股票：", stock_nums, "总资产：", round(capital + stock_nums * custommodule.get_stock_price(stock_code, stock_dates[-1]), 2), "收益率：", str(RoR) + "%")

        # 插入回测结果到backtest表
        with sqlite3.connect('db_file.db') as conn:
            cursor = conn.cursor()

            # 保存该策略的回测结果，不影响其他策略的回测结果
            if RoR != 0:
                cursor.execute(f"""
                                    INSERT INTO backtest (stock_code, {strategys[strategy]})
                                    VALUES (?, ?)
                                    ON CONFLICT(stock_code) DO UPDATE SET {strategys[strategy]} = excluded.{custommodule.Strategy.strategys()[strategy]}
                                """, (stock_code, RoR))
                conn.commit()

    backtest_result(strategy=strategys[strategy])

# 回测结果查询
def backtest_result(strategy=None):

    # 显示所有交易策略，供用户选择
    with sqlite3.connect('db_file.db') as conn:
        cursor = conn.cursor()
        cursor.execute("PRAGMA table_info(backtest)")
        cols = [col[1] for col in cursor.fetchall()][1:]
        strategys = dict(enumerate(cols, start=1))

    # 检查输入的交易策略参数，如果为None就让用户选择
    if strategy is None:
        # 用户选择显示的交易策略
        while True:
            print("当前交易策略有：")
            for x, y in strategys.items():
                print(f"{x}. {y}")
            strategy = input("请输入要查询的交易策略序号（默认为 1）：")

            if strategy and (int(strategy) in strategys.keys()):
                strategy = strategys[int(strategy)]
                break
            elif strategy == "":
                strategy = strategys[1]
                break
            else:
                print('❌ 输入有误，请重新输入！')

    result = []  # 用于存储策略回测的结果

    # 查询回测结果
    with (sqlite3.connect('db_file.db') as conn):
        sql = f"SELECT stock_code, {strategy} FROM backtest ORDER BY {strategy}"
        result = pd.read_sql(sql, conn).dropna()
        print(result)

    # 显示回测结果
    print(strategy, "回测结果：")
    if result[strategy].count() != 0:print("胜     率：", round((result[strategy] > 0).sum() / result[strategy].count() * 100, 2), "%")
    print("平均盈利率：", round(result[strategy].mean(), 2), "%")
    print("中位盈利率：", result[strategy].median(), "%")
    print("最大盈利率：", result[strategy].max(), "%")
    print("最小盈利率：", result[strategy].min(), "%")
    
    return result

def menu():
    while True:
        print("\n","====交易策略回测====")
        print("1、自选股回测")
        print("2、所有股回测")
        print("3、指定股票回测")
        print("4、查看回测结果")
        print("5、交易策略管理")
        print("0、返回主菜单")
        print("------------------")
        choice = input("请输入您的选择：")
        if choice not in ["0", "1", "2", "3", "4", "5"]:
            print("输入有误，请重新输入！")
        else:
            match choice:
                case "0":
                    break
                case "1":
                    backtest(stock_codes=[code[0] for code in custommodule.get_watchlist()])
                case "2":
                    backtest(stock_codes=[code[0] for code in custommodule.get_all_stocks_info()])
                case "3":
                    backtest(stock_codes=input("请输入股票代码（多个股票代码用,隔开）：").split(","))
                case "4":
                    backtest_result()
                case "5":
                    custommodule.Strategy.reset_strategys()
